package com.example.gateway.filter;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.data.redis.core.ReactiveStringRedisTemplate;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.List;
import java.util.Map;

@Component
public class RedisDynamicIpRateLimiterFilter extends AbstractGatewayFilterFactory<Object> {

    @Autowired
    private ReactiveStringRedisTemplate redisTemplate;

    private static final String PREFIX = "rate_limit:ip:";
    private static final String AUTO_BLOCK_KEY = "config:autoBlockThreshold";

    @Override
    public GatewayFilter apply(Object config) {
        return (exchange, chain) -> {
            String ip = getIpAddress(exchange);
            String path = exchange.getRequest().getPath().value();

            Mono<Boolean> isBlacklisted = redisTemplate.opsForSet().isMember("ip:blacklist", ip);
            Mono<Boolean> isWhitelisted = redisTemplate.opsForSet().isMember("ip:whitelist", ip);

            Mono<Map<Object, Object>> rateLimitParams = redisTemplate.opsForHash()
                    .entries(PREFIX + ip)
                    .collectMap(Map.Entry::getKey, Map.Entry::getValue);

            // 读取自动拉黑阈值
            Mono<String> autoBlockThresholdMono = redisTemplate.opsForValue().get(AUTO_BLOCK_KEY);

            return Mono.zip(isBlacklisted, isWhitelisted, rateLimitParams, autoBlockThresholdMono)
                    .flatMap(tuple -> {
                        boolean black = tuple.getT1();
                        boolean white = tuple.getT2();
                        Map<Object, Object> params = tuple.getT3();
                        int autoBlockThreshold = parseOrDefault(tuple.getT4(), 200);   // 默认 200 qps

                        if (black) {
                            exchange.getResponse().setStatusCode(HttpStatus.FORBIDDEN);
                            return exchange.getResponse().setComplete();
                        }

                        if (white) {
                            return chain.filter(exchange);
                        }

                        // 令牌桶参数（基础限流参数，没配置则默认）
                        int replenishRate = params.containsKey("replenishRate") ?
                                parseOrDefault(params.get("replenishRate"), 20) : 20;
                        int burstCapacity = params.containsKey("burstCapacity") ?
                                parseOrDefault(params.get("burstCapacity"), 30) : 30;

                        String tokenKey = PREFIX + ip + ":" + path + ":tokens";
                        String timestampKey = PREFIX + ip + ":" + path + ":timestamp";

                        long now = System.currentTimeMillis();

                        return redisTemplate.opsForValue().multiGet(List.of(tokenKey, timestampKey))
                                .flatMap(values -> {
                                    int tokens = burstCapacity;
                                    long lastRefill = now;

                                    if (values != null && values.size() == 2) {
                                        tokens = parseOrDefault(values.get(0), burstCapacity);
                                        lastRefill = parseOrDefault(values.get(1), now);
                                    }

                                    long delta = now - lastRefill;
                                    int newTokens = (int) (delta * replenishRate / 1000);
                                    tokens = Math.min(tokens + newTokens, burstCapacity);

                                    if (tokens > 0) {
                                        tokens--;
                                    } else {
                                        exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                                        return exchange.getResponse().setComplete();
                                    }

                                    Mono<Boolean> setTokens = redisTemplate.opsForValue().set(tokenKey, String.valueOf(tokens), Duration.ofSeconds(1));
                                    Mono<Boolean> setTimestamp = redisTemplate.opsForValue().set(timestampKey, String.valueOf(now), Duration.ofSeconds(1));

                                    String countKey = "rate_limit:count:" + ip + ":" + path;
                                    Mono<Long> incrementCount = redisTemplate.opsForValue().increment(countKey);
                                    redisTemplate.expire(countKey, Duration.ofSeconds(1)).subscribe();

                                    return Mono.zip(setTokens, setTimestamp, incrementCount)
                                            .flatMap(results -> {
                                                long count = results.getT3();

                                                if (count > autoBlockThreshold) {
                                                    redisTemplate.opsForSet().add("ip:blacklist", ip).subscribe();
                                                    exchange.getResponse().setStatusCode(HttpStatus.FORBIDDEN);
                                                    return exchange.getResponse().setComplete();
                                                }

                                                return chain.filter(exchange);
                                            });
                                });
                    });
        };
    }

    private int parseOrDefault(Object val, int defaultVal) {
        try {
            return val != null ? Integer.parseInt(val.toString()) : defaultVal;
        } catch (Exception e) {
            return defaultVal;
        }
    }

    private long parseOrDefault(Object val, long defaultVal) {
        try {
            return val != null ? Long.parseLong(val.toString()) : defaultVal;
        } catch (Exception e) {
            return defaultVal;
        }
    }

    private String getIpAddress(ServerWebExchange exchange) {
        String ip = exchange.getRequest().getHeaders().getFirst("X-Forwarded-For");
        if (ip == null || ip.isEmpty()) {
            ip = exchange.getRequest().getRemoteAddress().getAddress().getHostAddress();
        } else if (ip.contains(",")) {
            ip = ip.split(",")[0];
        }
        return ip;
    }
}
